package com.dman.multi.core;

import java.io.EOFException;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

public class NetworkStream {
	
	public static final long BUFFER_SIZE = 4096;
	private Selector selector = null;
	private Map<Section, SelectionKey>secKeyLookup = null;
	
	public NetworkStream() throws IOException {
		selector = Selector.open();
		secKeyLookup = new HashMap<Section, SelectionKey>();
	}
	
	/**
	 * This method registers a given section with a selector. The section is 
	 * registered for read operation only. This also attaches the section to the
	 * selection key returned from the selector.
	 * 
	 * @param Section object
	 * @return SelectionKey
	 * @throws IOException
	 */
	public void registerSection(Section sec) throws IOException{
		SocketChannel sockChannel = sec.call();
		SelectionKey selKey = sockChannel.register(selector, SelectionKey.OP_READ, sec);
		secKeyLookup.put(sec, selKey);
	}
	
	/**
	 * This method performs a select() operation on a selector and returns the set 
	 * of keys which are selected  for read operation.
	 * 
	 * @return Set of SelectionKey
	 * @throws IOException
	 * @throws InterruptedException 
	 */
	public Set<SelectionKey> getReadySockets(long timeout) throws IOException, InterruptedException {
		selector.select(timeout);
		Set<SelectionKey> setOfSock = selector.selectedKeys();
		return setOfSock;
	}
	
	/**
	 * This method reads a socket into a byte buffer and returns the same. It clears
	 * the buffer before it populates data into it. It also flips the same after
	 * completion of read, so that it can readily be used for writing.
	 * 
	 * @param SelectionKey
	 * @return a ByteBuffer
	 * @throws IOException
	 */
	public ByteBuffer readSocket(SelectionKey selKey) throws EOFException, IOException{
		Section sec = (Section) selKey.attachment();
		ByteBuffer bb = sec.read();
		return bb;
	}
	
	/**
	 * This method re-initializes a section, this is to help out when a socket
	 * is hung, or not responding.
	 * 
	 * @param Section
	 * @throws IOException 
	 */
	public void reinit(Section sec) throws IOException{
		SelectionKey selKey = secKeyLookup.get(sec);
		selKey.cancel();
		sec.close();
		registerSection(sec);
	}
	
	/**
	 * This method closes the sockets if the download is interrupted. This also takes
	 * care of sockets which might already be closed, by the interrupt.
	 * 
	 * @throws IOException
	 */
	public void close() throws IOException{
		for(SelectionKey selKey: selector.keys()){
			SocketChannel sockChannel = (SocketChannel) selKey.channel();
			if(sockChannel.isOpen())
				sockChannel.close();
		}
		selector.close();
	}
}
